{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "981813f9",
   "metadata": {},
   "source": [
    "# A2A Client example\n",
    "\n",
    "In order to run this example, you need to run the server:\n",
    "\n",
    "```\n",
    "python a2a_min/examples/base/server_example.py\n",
    "```\n",
    "\n",
    "or using the utility script:\n",
    "```\n",
    "a2a_min/scripts/run_server_example_base.sh\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8186417",
   "metadata": {},
   "source": [
    "### Initialize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8de57ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "from pathlib import Path\n",
    "from uuid import uuid4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f36351f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add the project's root directory to the Python path\n",
    "import sys\n",
    "\n",
    "PROJECT_ROOT = Path.cwd().resolve().parent.parent.parent\n",
    "print(f\"PROJECT_ROOT = {PROJECT_ROOT}\")\n",
    "if PROJECT_ROOT not in sys.path:\n",
    "    sys.path.append(str(PROJECT_ROOT))\n",
    "\n",
    "print(f\"sys.path: {sys.path}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "698138e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from a2a_min.base.client.card_resolver import A2ACardResolver\n",
    "from a2a_min.base.client.client import A2AClient\n",
    "from a2a_min.base.types import Message, TaskSendParams, TaskState, TextPart"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68242ede",
   "metadata": {},
   "outputs": [],
   "source": [
    "A2A_SERVER_URL = \"http://localhost:10000\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a66a14e",
   "metadata": {},
   "source": [
    "### Get A2A server's card"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9087a055",
   "metadata": {},
   "outputs": [],
   "source": [
    "card_resolver = A2ACardResolver(A2A_SERVER_URL)\n",
    "card = card_resolver.get_agent_card()\n",
    "\n",
    "print(card.model_dump_json(exclude_none=True, indent=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7baedb1c",
   "metadata": {},
   "source": [
    "### Create an A2A client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36b7afca",
   "metadata": {},
   "outputs": [],
   "source": [
    "client = A2AClient(agent_card=card)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9923154",
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_json(data):\n",
    "    if hasattr(data, \"model_dump_json\"):\n",
    "        return data.model_dump_json(indent=2, exclude_none=True)\n",
    "    return data.json(indent=2, exclude_none=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d87825d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def completeTask(a2a_client: A2AClient, prompt: str, taskId, sessionId):\n",
    "    \"\"\" Run a task with the given prompt \"\"\"\n",
    "    if prompt == \":q\" or prompt == \"quit\":\n",
    "        return False\n",
    "\n",
    "    message = Message(role=\"user\", parts=[TextPart(text=prompt)])\n",
    "    tasK_send_params = TaskSendParams(id=taskId, sessionId=sessionId, acceptedOutputModes=[\"text\"], message=message)\n",
    "\n",
    "    taskResult = await a2a_client.send_task(tasK_send_params)\n",
    "    print(f\"\\n{to_json(taskResult)}\")\n",
    "\n",
    "    ## if the result is that more input is required, loop again.\n",
    "    state = TaskState(taskResult.result.status.state)\n",
    "    if state.name == TaskState.INPUT_REQUIRED.name:\n",
    "        return await completeTask(a2a_client, prompt, taskId, sessionId)\n",
    "    else:\n",
    "        return False    # We are done with this task.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6ae8bea",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"What do you want to send to the agent? (:q or quit to exit)\"\n",
    "\n",
    "continue_loop = True\n",
    "sessionId = uuid4().hex\n",
    "while continue_loop:\n",
    "    taskId = uuid4().hex\n",
    "    print(\"=========  starting a new task ======== \")\n",
    "    continue_loop = await completeTask(client, prompt, taskId, sessionId)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
